# coding:utf-8
from numpy.core.numeric import zeros, arange, concatenate, ones, array
from numpy.ma.core import arctan2, log, exp, sqrt, conjugate, size
from numpy.fft.fftpack import fft, ifft
from scipy.signal.signaltools import triang

class Tfa(object):
    """時間周波数分析による残響抑圧
    """
    
    def __init__(self, x, fs, FrameLength):
        self.x = x
        self.fs = fs
        self.FrameLength = FrameLength
        self.FrameNumber = size(x) / self.FrameLength
        self.Amp = zeros((self.FrameNumber, self.FrameLength / 2))
        self.Phase = zeros((self.FrameNumber, self.FrameLength / 2))
    
    def calc(self):
        """メインルーチン
        """
        for i in xrange(self.FrameNumber):
            xFrame = self.x[self.FrameLength * i / 2: self.FrameLength * i / 2 + self.FrameLength]
            XFrame = fft(xFrame)
            XFrame = XFrame[:size(XFrame, 0) / 2]
            self.Amp[i] = abs(XFrame)
            self.Phase[i] = arctan2(XFrame.imag, XFrame.real)
        Tr = 2 #残響時間
        t = self.FrameLength * arange(self.FrameNumber) / (2 * self.fs)
        E = exp(-log(10 ** 6) * t / Tr)
        for k in xrange(self.FrameLength / 2):
            for i in xrange(1, self.FrameNumber):
                ReverbePower = 0
                for ii in xrange(i):
                    ReverbePower += self.Amp[ii][k] ** 2 * E[i - ii]
                DryPower = self.Amp[i][k] ** 2 - ReverbePower
                if DryPower < 0:
                    DryPower = 0
                self.Amp[i][k] = sqrt(DryPower)
        y = zeros(size(self.x, 0)) #残響が除去された信号
        for i in xrange(self.FrameNumber):
            YFrame = self.Amp[i] * exp(1j * self.Phase[i])
            YFrame = concatenate((YFrame, conjugate(array(list(reversed(YFrame))))))
            yFrame = ifft(YFrame)
            Window = triang(self.FrameLength)
            if i == 0:
                startWindow = Window[:]
                startWindow[: self.FrameLength / 2] = ones(self.FrameLength / 2)
                y[self.FrameLength * i / 2: self.FrameLength * i / 2 + self.FrameLength] += yFrame * startWindow  
            else:
                y[self.FrameLength * i / 2: self.FrameLength * i / 2 + self.FrameLength] += yFrame * Window
        return y